package db

import (
	"context"
	"database/sql"
	"fmt"
	pgdb "originforce.club/bilin-mock/internal/db/sqlc"
)

// TxRunner run transaction container
type TxRunner struct {
	*pgdb.Queries
	db *sql.DB
}

// NewTxRunner create a new TxRunner
func NewTxRunner(db *sql.DB) *TxRunner {
	return &TxRunner{
		db:      db,
		Queries: pgdb.New(db),
	}
}

// ExecTx 事务执行
func (txRunner *TxRunner) ExecTx(ctx context.Context, fn func(*sql.Tx) error) error {
	tx, err := txRunner.db.BeginTx(ctx, nil)
	if err != nil {
		return err
	}
	err = fn(tx)
	if err != nil {
		if rbErr := tx.Rollback(); rbErr != nil {
			return fmt.Errorf("tx err:%v, rollback err:%v", err, rbErr)
		}
		return err
	}

	return tx.Commit()
}

// SqlcExecTx 事务执行sqlcQuery
func (txRunner *TxRunner) SqlcExecTx(ctx context.Context, fn func(*pgdb.Queries) error) error {
	tx, err := txRunner.db.BeginTx(ctx, nil)
	if err != nil {
		return err
	}
	txq := pgdb.New(tx)
	err = fn(txq)
	if err != nil {
		if rbErr := tx.Rollback(); rbErr != nil {
			return fmt.Errorf("tx err:%v, rollback err:%v", err, rbErr)
		}
		return err
	}

	return tx.Commit()
}

// StmtExecTx 事务执行rawStatement
func (txRunner *TxRunner) StmtExecTx(ctx context.Context, query string, fn func(*sql.Stmt) error) error {
	tx, err := txRunner.db.BeginTx(ctx, nil)
	if err != nil {
		return err
	}
	stmt, err := tx.PrepareContext(ctx, query)
	if err != nil {
		return err
	}
	err = fn(stmt)
	if err != nil {
		if rbErr := tx.Rollback(); rbErr != nil {
			return fmt.Errorf("tx err:%v, rollback err:%v", err, rbErr)
		}
		return err
	}

	return tx.Commit()
}
